from typing import Any

import torch

import torchair
from torchair import register_fx_node_ge_converter
from torchair.ge import Tensor

@register_fx_node_ge_converter(torch.ops.npu.fused_add_topk_div.default)
def convert_fused_add_topk_div(x: Tensor, add_num: Tensor, mapping_num:Tensor, mapping_table: Tensor,  group_num: int,
                               group_topk: int, n: int, k: int, activate_type: int,
                               is_norm: bool, enableExpertMapping: bool, scale: float, y: Tensor = None,
                               indices: Tensor = None, meta_outputs: Any = None):
    return torchair.ge.custom_op(
        "FusedAddTopkDiv",
        inputs={
            "x": x,
            "add_num": add_num,
            "mapping_num": mapping_num,
            "mapping_table": mapping_table
        },
        attrs={
            "group_num": torchair.ge.attr.Int(group_num),
            "group_topk": torchair.ge.attr.Int(group_topk),
            "n": torchair.ge.attr.Int(n),
            "k": torchair.ge.attr.Int(k),
            "activate_type": torchair.ge.attr.Int(activate_type),
            "is_norm": torchair.ge.attr.Bool(is_norm),
            "enableExpertMapping": torchair.ge.attr.Bool(enableExpertMapping),
            "scale": torchair.ge.attr.Float(scale)
        },
        outputs=['y', 'indices']
    )